#!/bin/bash

# NVIDIA NeMo Agent Toolkit AI对话机器人 - 一键安装脚本
# 支持 Linux/macOS/Windows(WSL)

set -e

echo "🚀 智能工地巡检助理V0.5 - 一键安装"
echo "=================================================="

# 检测操作系统
detect_os() {
    if [[ "$OSTYPE" == "linux-gnu"* ]]; then
        echo "检测到操作系统: Linux"
        OS="linux"
    elif [[ "$OSTYPE" == "darwin"* ]]; then
        echo "检测到操作系统: macOS"
        OS="macos"
    elif [[ "$OSTYPE" == "msys" ]] || [[ "$OSTYPE" == "cygwin" ]]; then
        echo "检测到操作系统: Windows"
        OS="windows"
    else
        echo "⚠️  未知操作系统: $OSTYPE"
        echo "请在 Linux、macOS 或 Windows WSL 环境中运行此脚本"
        exit 1
    fi
}

# 检查必要工具
check_requirements() {
    echo "📋 检查系统要求..."
    
    # 检查Python
    if ! command -v python3 &> /dev/null; then
        echo "❌ Python 3 未安装"
        echo "请先安装 Python 3.12 或更高版本"
        exit 1
    fi
    
    python_version=$(python3 --version | cut -d' ' -f2 | cut -d'.' -f1-2)
    echo "✅ Python 版本: $python_version"
    
    # 检查Node.js
    if ! command -v node &> /dev/null; then
        echo "❌ Node.js 未安装"
        echo "请先安装 Node.js 18 或更高版本"
        exit 1
    fi
    
    node_version=$(node --version)
    echo "✅ Node.js 版本: $node_version"
    
    # 检查Git
    if ! command -v git &> /dev/null; then
        echo "❌ Git 未安装"
        echo "请先安装 Git"
        exit 1
    fi
    
    echo "✅ Git 已安装"
}

# 安装uv包管理器
install_uv() {
    echo "📦 安装 uv 包管理器..."
    
    if ! command -v uv &> /dev/null; then
        echo "正在下载并安装 uv..."
        curl -LsSf https://astral.sh/uv/install.sh | sh
        
        # 添加到PATH
        export PATH="$HOME/.local/bin:$PATH"
        
        if ! command -v uv &> /dev/null; then
            echo "❌ uv 安装失败"
            exit 1
        fi
    fi
    
    echo "✅ uv 包管理器已安装"
}

# 克隆和设置项目
setup_project() {
    echo "📥 设置项目..."
    
    # 记录当前目录
    PROJECT_ROOT=$(pwd)
    
    # 如果目录不存在，克隆项目
    if [ ! -d "NeMo-Agent-Toolkit" ]; then
        echo "正在克隆 NVIDIA NeMo Agent Toolkit..."
        git clone https://github.com/NVIDIA/NeMo-Agent-Toolkit.git
        
        echo "正在初始化子模块..."
        cd NeMo-Agent-Toolkit
        git submodule update --init --recursive
        cd "$PROJECT_ROOT"
    else
        echo "项目目录已存在..."
    fi
    
    # 进入NeMo-Agent-Toolkit目录进行Python环境设置
    cd NeMo-Agent-Toolkit
    
    # 创建Python虚拟环境
    echo "正在创建Python虚拟环境..."
    uv venv --seed .venv --python 3.12
    
    # 激活虚拟环境并安装依赖
    echo "正在安装Python依赖..."
    source .venv/bin/activate
    cp -rf ../hc_test examples
    uv pip install -e .
    uv pip install -e '.[langchain]'
    uv pip install tavily-python
    uv pip install 'httpx[socks]'
    uv pip install -e examples/hc_test/user_report
    
    echo "✅ 后端依赖安装完成"
    
    # 返回项目根目录
    cd "$PROJECT_ROOT"
}
install_mysql() {
    PROJECT_ROOT=$(pwd)
    cd NeMo-Agent-Toolkit
    docker compose -f examples/deploy/docker-compose.mysql.yml up -d
    source .venv/bin/activate
    nat object-store mysql --host 127.0.0.1 --username root --password ${MYSQL_ROOT_PASSWORD} --port 3306 my-bucket upload ./examples/object_store/user_report/data/object_store/
    cd "$PROJECT_ROOT"
    echo "✅ MySQL 已安装"
}
# 设置前端
setup_frontend() {
    echo "🎨 设置前端..."
    
    # 检查前端目录是否存在
    if [ ! -d "external/aiqtoolkit-opensource-ui" ]; then
        echo "❌ 前端目录不存在: external/aiqtoolkit-opensource-ui"
        echo "请确保子模块已正确初始化"
        exit 1
    fi
    
    # 进入前端目录
    cd external/aiqtoolkit-opensource-ui
    
    echo "正在安装前端依赖..."
    npm install
    
    echo "✅ 前端依赖安装完成"
    
    # 返回项目根目录
    cd ../..
}

# 创建配置文件
create_config() {
    echo "⚙️  创建配置文件..."
    
    # 确保在项目根目录
    PROJECT_ROOT=$(pwd)
    
    # 在NeMo-Agent-Toolkit目录中创建配置文件
    cd NeMo-Agent-Toolkit
    mkdir -p configs
    
    cat > configs/hackathon_config.yml << 'EOF'
# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

general:
  front_end:
    _type: fastapi
    object_store: report_object_store
    cors:
      allow_origins: ['*']

object_stores:
  report_object_store:
    _type: mysql
    host: localhost
    port: 3306
    username: root
    password: 
    bucket_name: my-bucket

llms:
  # 默认使用Bailian API (用户可修改)
  default_llm:
    _type: openai
    model_name: "qwen-plus"
    api_key: ""
    base_url: "https://dashscope.aliyuncs.com/compatible-mode/v1"
    temperature: 0.7
    max_tokens: 2048

functions:
  get_weld_user_report:
    _type: get_weld_user_report
    object_store: report_object_store
    description: >
      从对象存储桶中读取指定用户的焊接报告。
      参数:
        user_id: str: 用户ID(例如 P12345、P24680、P67890)。
        date: str | null: 报告日期，格式 YYYY-MM-DD。
          如果未提供日期，则默认获取该用户的最新报告。

  put_weld_user_report:
    _type: put_weld_user_report
    object_store: report_object_store
    description: >
      将新的焊接报告写入对象存储桶。
      如果该用户在指定日期的报告已经存在，则写入失败，
      并提示用户该报告已存在，不会覆盖或删除已有数据。
      参数:
        report: str: 要写入的报告内容(JSON 字符串)。
        user_id: str: 用户ID。
        date: str | null: 报告日期，格式 YYYY-MM-DD。
          如果未提供日期，则默认存储为 "latest"。

  update_weld_user_report:
    _type: update_weld_user_report
    object_store: report_object_store
    description: >
      更新指定用户的焊接报告。
      如果该用户在指定日期的报告不存在，则会新建一份。
      不会删除已有的数据，只会更新或新增。
      参数:
        report: str: 要更新的报告内容(JSON 字符串)。
        user_id: str: 用户ID。
        date: str | null: 报告日期，格式 YYYY-MM-DD。
          如果未提供日期，则默认存储为 "latest"。

  delete_weld_user_report:
    _type: delete_weld_user_report
    object_store: report_object_store
    description: >
      删除指定用户的焊接报告。
      如果报告不存在，则会报错提示。
      注意：只有在明确需要删除的情况下才运行此功能，
      插入失败时也不能用删除+重建的方式替代。
      参数:
        user_id: str: 用户ID。
        date: str | null: 报告日期，格式 YYYY-MM-DD。
          如果未提供日期，则默认删除 "latest"。

workflow:
  _type: react_agent
  tool_names: [get_weld_user_report, put_weld_user_report, update_weld_user_report, delete_weld_user_report]
  llm_name: default_llm
  verbose: true
  handle_parsing_errors: true
  max_retries: 2
EOF
    
    echo "✅ 配置文件创建完成"
    
    # 返回项目根目录
    cd "$PROJECT_ROOT"
}

# 创建启动脚本
create_scripts() {
    echo "📝 创建启动脚本..."
    
    # 确保在项目根目录
    PROJECT_ROOT=$(pwd)
    
    # 在NeMo-Agent-Toolkit目录中创建启动脚本
    cd NeMo-Agent-Toolkit
    
    # 创建启动脚本
    cat > start.sh << 'EOF'
#!/bin/bash

echo "🚀 启动 NVIDIA NeMo Agent Toolkit AI对话机器人"
echo "=============================================="

# 获取项目根目录和NeMo目录
NEMO_DIR=$(pwd)
PROJECT_ROOT=$(dirname "$NEMO_DIR")

# 设置环境变量
export TAVILY_API_KEY=

# 激活Python虚拟环境
source .venv/bin/activate

# 启动后端服务
echo "📡 启动后端服务..."
aiq serve --config_file configs/hackathon_config.yml --host 0.0.0.0 --port 8001 &
BACKEND_PID=$!

# 等待后端启动
echo "⏳ 等待后端服务启动..."
sleep 10

# 启动前端服务
echo "🎨 启动前端服务..."
cd "$PROJECT_ROOT/external/aiqtoolkit-opensource-ui"
npm run dev &
FRONTEND_PID=$!

# 返回NeMo目录
cd "$NEMO_DIR"

echo ""
echo "✅ 系统启动完成！"
echo ""
echo "🌐 访问地址:"
echo "   前端界面: http://localhost:3000"
echo ""
echo "📝 测试建议:"
echo "   巡检报告查询: '操作员P003的最新巡检报告？'"
echo ""
echo "🛑 停止服务: 按 Ctrl+C 或运行 ./stop.sh"
echo ""

# 保存进程ID
echo $BACKEND_PID > .backend.pid
echo $FRONTEND_PID > .frontend.pid

# 等待用户中断
wait
EOF

    # 创建停止脚本
    cat > stop.sh << 'EOF'
#!/bin/bash

echo "🛑 停止 NVIDIA NeMo Agent Toolkit AI对话机器人"
echo "=============================================="

# 停止后端服务
if [ -f .backend.pid ]; then
    BACKEND_PID=$(cat .backend.pid)
    if ps -p $BACKEND_PID > /dev/null; then
        echo "停止后端服务 (PID: $BACKEND_PID)..."
        kill $BACKEND_PID
    fi
    rm -f .backend.pid
fi

# 停止前端服务
if [ -f .frontend.pid ]; then
    FRONTEND_PID=$(cat .frontend.pid)
    if ps -p $FRONTEND_PID > /dev/null; then
        echo "停止前端服务 (PID: $FRONTEND_PID)..."
        kill $FRONTEND_PID
    fi
    rm -f .frontend.pid
fi

# 清理其他相关进程
pkill -f "aiq serve" 2>/dev/null || true
pkill -f "next dev" 2>/dev/null || true

echo "✅ 所有服务已停止"
EOF

    # 添加执行权限
    chmod +x start.sh stop.sh
    
    echo "✅ 启动脚本创建完成"
    
    # 返回项目根目录
    cd "$PROJECT_ROOT"
}

# 主安装流程
main() {
    detect_os
    check_requirements
    install_uv
    setup_project
    install_mysql
    setup_frontend
    create_config
    create_scripts
    
    echo ""
    echo "🎉 安装完成！"
    echo "=============="
    echo ""
    echo "📁 项目根目录: $(pwd)"
    echo "📁 NeMo项目目录: $(pwd)/NeMo-Agent-Toolkit"
    echo ""
    echo "🚀 快速启动:"
    echo "   cd NeMo-Agent-Toolkit && ./start.sh"
    echo ""
    echo "🛑 停止服务:"
    echo "   cd NeMo-Agent-Toolkit && ./stop.sh"
    echo ""
    echo "⚙️  自定义配置:"
    echo "   编辑 NeMo-Agent-Toolkit/configs/hackathon_config.yml 文件"
    echo "   可修改 API密钥、模型名称、base_url 等"
    echo ""
    echo "📚 更多信息:"
    echo "   https://github.com/NVIDIA/NeMo-Agent-Toolkit"
    echo ""
}

# 运行主函数
main "$@"

